import numpy as np
from collections import Counter
import matplotlib.pyplot as plt
from sklearn.datasets import make_classification

class kNN:
    def __init__(self, k=3):
        self.k = k
        self.X_train = None
        self.y_train = None
        
    def fit(self, X, y):
        """训练kNN模型（只是存储数据）"""
        self.X_train = X
        self.y_train = y
        return self
    
    def predict(self, X):
        """预测新样本的类别"""
        predictions = [self._predict(x) for x in X]
        return np.array(predictions)
    
    def _predict(self, x):
        """预测单个样本的类别"""
        # 计算距离
        distances = np.linalg.norm(self.X_train - x, axis=1)
        
        # 获取最近的k个邻居的索引
        k_indices = np.argsort(distances)[:self.k]
        
        # 获取k个邻居的标签
        k_nearest_labels = self.y_train[k_indices]
        
        # 投票决定类别
        most_common = Counter(k_nearest_labels).most_common(1)
        return most_common[0][0]

# 测试示例
if __name__ == "__main__":
    # 生成分类数据
    X, y = make_classification(n_samples=200, n_features=2, n_redundant=0, 
                              n_informative=2, n_clusters_per_class=1, 
                              random_state=42)
    
    # 划分训练测试集
    split = 150
    X_train, X_test = X[:split], X[split:]
    y_train, y_test = y[:split], y[split:]
    
    # 训练kNN模型
    knn = kNN(k=5)
    knn.fit(X_train, y_train)
    
    # 预测
    y_pred = knn.predict(X_test)
    
    # 计算准确率
    accuracy = np.mean(y_pred == y_test)
    print(f"测试集准确率: {accuracy:.3f}")
    
    # 可视化结果
    plt.figure(figsize=(12, 5))
    
    # 训练数据
    plt.subplot(131)
    plt.scatter(X_train[:, 0], X_train[:, 1], c=y_train, cmap='viridis', alpha=0.7)
    plt.title("训练数据")
    
    # 测试数据真实标签
    plt.subplot(132)
    plt.scatter(X_test[:, 0], X_test[:, 1], c=y_test, cmap='viridis', alpha=0.7)
    plt.title("测试数据真实标签")
    
    # 测试数据预测标签
    plt.subplot(133)
    plt.scatter(X_test[:, 0], X_test[:, 1], c=y_pred, cmap='viridis', alpha=0.7)
    plt.title("测试数据预测标签")
    
    plt.tight_layout()
    plt.show()
    
    # 演示单个预测过程
    print("\n单个样本预测演示:")
    test_sample = X_test[0]
    print(f"测试样本: {test_sample}")
    print(f"真实标签: {y_test[0]}")
    print(f"预测标签: {y_pred[0]}")